# -*- coding:utf-8 -*-

import torch
import torch.nn as nn


class ConvBN(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 stride=1,
                 padding=1,
                 dilations=1,
                 has_relu=True):
        super(ConvBN, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(
                in_channels,
                out_channels,
                kernel_size,
                stride,
                padding,
                dilations,
                bias=False
            ),
            nn.BatchNorm2d(out_channels)
        )
        self.has_relu = has_relu
        if self.has_relu:
            self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        if self.has_relu:
            x = self.relu(x)
        return x


class Concat(nn.Module):
    def __init__(self, dim=1):
        super(Concat, self).__init__()
        self.dim = dim

    def forward(self, inputs):
        return torch.cat(inputs, dim=self.dim)


class OSABlock(nn.Module):
    def __init__(self,
                 in_channels,
                 mid_channels,
                 out_channels,
                 stride=2,
                 repeat=5):
        super(OSABlock, self).__init__()
        assert stride in [1, 2]
        assert repeat >= 2
        self.repeat = repeat

        self.conv1 = ConvBN(in_channels, mid_channels, stride=stride)

        for i in range(repeat - 1):
            self._modules['conv{}'.format(i + 2)] = ConvBN(mid_channels, mid_channels)

        self.concat = Concat()
        self.reduce = ConvBN(mid_channels * repeat, out_channels, kernel_size=1, padding=0)

    def forward(self, x):
        layers = []
        for i in range(self.repeat):
            x = self._modules['conv{}'.format(i + 1)](x)
            layers.append(x)
        x = self.concat(layers)
        x = self.reduce(x)
        return x


class PFLDInference(nn.Module):
    def __init__(self):
        super(PFLDInference, self).__init__()

        self.stem1 = ConvBN(3, 32, stride=2)

        self.osa2 = OSABlock(32, 32, 64, stride=2, repeat=5)
        self.osa3 = OSABlock(64, 64, 96, stride=2, repeat=5)
        self.osa4 = OSABlock(96, 96, 144, stride=2, repeat=5)
        self.osa5 = OSABlock(144, 128, 192, stride=2, repeat=5)

        # landmarks
        self.kp_conv3 = ConvBN(96, 64)
        self.kp_pool3 = nn.AdaptiveAvgPool2d((1, 1))
        self.kp_conv4 = ConvBN(144, 64)
        self.kp_pool4 = nn.AdaptiveAvgPool2d((1, 1))
        self.kp_conv5 = ConvBN(192, 128)
        self.kp_pool5 = nn.AdaptiveAvgPool2d((1, 1))
        self.kp_concat = Concat()
        self.kp_loc_fc = nn.Conv2d(256, 256, kernel_size=1, padding=0)
        self.kp_loc = nn.Conv2d(256, 196, kernel_size=1, padding=0)

    def forward(self, x):  # x: 3, 256, 256
        stem1 = self.stem1(x)
        osa2 = self.osa2(stem1)
        osa3 = self.osa3(osa2)
        osa4 = self.osa4(osa3)
        osa5 = self.osa5(osa4)

        kp3 = self.kp_pool3(self.kp_conv3(osa3))
        kp4 = self.kp_pool4(self.kp_conv4(osa4))
        kp5 = self.kp_pool5(self.kp_conv5(osa5))

        kp_concat = self.kp_concat([kp3, kp4, kp5])
        kp_loc = self.kp_loc(self.kp_loc_fc(kp_concat))

        return osa3, kp_loc.view(kp_loc.size(0), -1)


class AuxiliaryNet(nn.Module):
    def __init__(self):
        super(AuxiliaryNet, self).__init__()

        # auxiliary euler angles
        self.aux4 = OSABlock(96, 96, 144, stride=2, repeat=3)
        self.aux5 = OSABlock(144, 128, 192, stride=2, repeat=3)
        self.aux_pool = nn.AdaptiveMaxPool2d((1, 1))
        self.aux_fc1 = nn.Conv2d(192, 32, kernel_size=1, padding=0)
        self.aux_fc2 = nn.Conv2d(32, 3, kernel_size=1, padding=0)

    def forward(self, x):
        aux4 = self.aux4(x)
        aux5 = self.aux5(aux4)
        aux_pool = self.aux_pool(aux5)
        aux_fc = self.aux_fc2(self.aux_fc1(aux_pool))

        return aux_fc.view(aux_fc.size(0), -1)

# if __name__ == '__main__':
#     input = torch.randn(1, 3, 256, 256)
#     plfd_vovnet = PFLDVoVNet(with_auxiliary_net=True)
#     plfd_vovnet_eval = PFLDVoVNet(with_auxiliary_net=False)
#     print(plfd_vovnet)
#     print(plfd_vovnet_eval)
#     landmarks, errs = plfd_vovnet_eval(input)
#
#     print("landmarks.shape: {}".format(landmarks.shape[1:]))

